"""
MCP 基础数学运算服务器
提供核心数学运算功能：
- 基本运算：加减乘除
- 矩阵运算：4阶以下行列式计算
使用stdio传输协议与客户端通信
"""

from mcp.server.fastmcp import FastMCP
from typing import Union, List

# 创建MCP服务器实例
mcp = FastMCP("mcp-basic-math")

# ============ 基础四则运算 ============
@mcp.tool()
def add(a: float, b: float) -> float:
    """加法运算: 返回 a + b"""
    return a + b

@mcp.tool()
def subtract(a: float, b: float) -> float:
    """减法运算: 返回 a - b"""
    return a - b

@mcp.tool()
def multiply(a: float, b: float) -> float:
    """乘法运算: 返回 a * b"""
    return a * b

@mcp.tool()
def divide(a: float, b: float) -> Union[float, str]:
    """除法运算: 返回 a / b (处理除零错误)"""
    if b == 0:
        return "错误：除数不能为零"
    return a / b

# ============ 行列式计算 ============
@mcp.tool()
def calculate_determinant(matrix: List[List[float]]) -> Union[float, str]:
    """
    计算4阶以下方阵的行列式
    支持1×1到4×4的方阵
    """
    # 验证输入是否为方阵
    n = len(matrix)
    if n == 0 or any(len(row) != n for row in matrix):
        return "错误：输入必须是方阵"
    
    if n > 4:
        return "错误：只支持4阶及以下的行列式计算"
    
    # 1阶矩阵
    if n == 1:
        return matrix[0][0]
    
    # 2阶矩阵
    if n == 2:
        return matrix[0][0] * matrix[1][1] - matrix[0][1] * matrix[1][0]
    
    # 3阶矩阵（Sarrus法则）
    if n == 3:
        a, b, c = matrix[0]
        d, e, f = matrix[1]
        g, h, i = matrix[2]
        return (a*e*i + b*f*g + c*d*h) - (c*e*g + b*d*i + a*f*h)
    
    # 4阶矩阵（Laplace展开）
    det = 0.0
    for col in range(4):
        # 创建余子式（去掉第一行和第col列）
        minor = [row[:col] + row[col+1:] for row in matrix[1:]]
        minor_det = calculate_determinant(minor)
        if isinstance(minor_det, str):  # 如果出错
            return minor_det
        det += ((-1) ** col) * matrix[0][col] * minor_det
    
    return det

# ============ 服务器信息 ============
@mcp.tool()
def get_server_info() -> dict:
    """获取服务器基本信息"""
    return {
        "server_name": "mcp-basic-math",
        "version": "1.0",
        "description": "基础数学运算服务器（四则运算+行列式计算）",
        "transport": "stdio",
        "available_tools": ["add", "subtract", "multiply", "divide", "calculate_determinant"],
        "supported_operations": [
            "加法 (+)", "减法 (-)", "乘法 (*)", "除法 (/)", 
            "4阶以下行列式计算"
        ]
    }
# 在原有代码基础上添加以下内容：

@mcp.tool()
def matrix_multiply(matrix_a: List[List[float]], matrix_b: List[List[float]]) -> Union[List[List[float]], str]:
    """
    执行4阶以下矩阵乘法
    
    Args:
        matrix_a: 第一个矩阵 (m×n)
        matrix_b: 第二个矩阵 (n×p)
    
    Returns:
        乘积矩阵 (m×p)，如果维度不匹配则返回错误信息
    """
    # 验证矩阵A的列数是否等于矩阵B的行数
    rows_a = len(matrix_a)
    if rows_a == 0:
        return "错误：矩阵A不能为空"
    
    cols_a = len(matrix_a[0])
    rows_b = len(matrix_b)
    
    if rows_b == 0:
        return "错误：矩阵B不能为空"
    
    cols_b = len(matrix_b[0])
    
    if cols_a != rows_b:
        return f"错误：矩阵A的列数({cols_a})必须等于矩阵B的行数({rows_b})"
    
    # 检查矩阵阶数限制
    if rows_a > 4 or cols_a > 4 or rows_b > 4 or cols_b > 4:
        return "错误：只支持4阶及以下的矩阵乘法"
    
    # 初始化结果矩阵
    result = [[0.0 for _ in range(cols_b)] for _ in range(rows_a)]
    
    # 执行矩阵乘法
    for i in range(rows_a):
        for j in range(cols_b):
            for k in range(cols_a):
                result[i][j] += matrix_a[i][k] * matrix_b[k][j]
    
    return result

# 更新服务器信息函数
@mcp.tool()
def get_server_info() -> dict:
    """获取服务器基本信息"""
    return {
        "server_name": "mcp-basic-math",
        "version": "1.1",  # 更新版本号
        "description": "基础数学运算服务器（四则运算+矩阵计算）",
        "transport": "stdio",
        "available_tools": [
            "add", "subtract", "multiply", "divide", 
            "calculate_determinant", "matrix_multiply"
        ],
        "supported_operations": [
            "加法 (+)", "减法 (-)", "乘法 (*)", "除法 (/)", 
            "4阶以下行列式计算",
            "4阶以下矩阵乘法"
        ]
    }
if __name__ == "__main__":
    # 启动MCP服务器
    mcp.run(transport="stdio")